# PA_STABLIZE.py: 稳定控制模块，用于通过IMU传感器和PID控制稳定姿态
import PA_IMU
import PA_AVGFILT
import array
from machine import I2C, Pin


# 中间变量定义
q = []
Sta_Pitch = 0
Sta_Roll = 0
kp_sta = 0.08
kd_sta = 0
p_origin = 0
r_origin = 0
time_p = 0
filter_data_p = 0
filter_data_r = 0
gyro_cal_sta = 0
gyro_x_fitted = 0


# 滑动平均滤波
# Pitch
gyro_data_p = array.array('i', [0] * 5)
f_gyro_data_p = PA_AVGFILT.avg_filiter(gyro_data_p)
# Roll
gyro_data_r = array.array('i', [0] * 5)
f_gyro_data_r = PA_AVGFILT.avg_filiter(gyro_data_r)
# 滑动平均滤波(角速度)
gyro_x_fit = array.array('i', [0] * 5)
gyro_x_fitting = PA_AVGFILT.avg_filiter(gyro_x_fit)


def get_imu_value():
    global q, filter_data_p, filter_data_r, time_p, p_origin, r_origin, gyro_cal_sta, gyro_x_fitted
    ay = PA_IMU.acc.get_values()

    if time_p <= 50:
        try:
            a = PA_IMU.IMUupdate(ay["GyX"], ay["GyY"], ay["GyZ"], ay["AcX"], ay["AcY"], ay["AcZ"])
            p_origin = round(a[0])
            r_origin = round(a[1])
            time_p = time_p + 1
            gyro_cal_sta = 0
        except:
            gyro_cal_sta = 0
    else:
        try:
            a = PA_IMU.IMUupdate(ay["GyX"], ay["GyY"], ay["GyZ"], ay["AcX"], ay["AcY"], ay["AcZ"])
            filter_data_p = f_gyro_data_p.avg(round(a[0]))
            filter_data_r = f_gyro_data_r.avg(round(a[1]))
            gyro_x_fitted = gyro_x_fitting.avg(round(ay["GyX"] / 65.5 * 0.0174533))  # 输出角速度
            gyro_cal_sta = 1
        except:
            gyro_cal_sta = 1

    ay = PA_IMU.acc.get_values()
    q = PA_IMU.IMUupdate(ay["GyX"], ay["GyY"], ay["GyZ"], ay["AcX"], ay["AcY"], ay["AcZ"])


def stab(bac_statue_p_or_pr):
    """
    该函数用于根据IMU传感器的数据调整机器狗的俯仰角（Pitch）和横滚角（Roll），
    并根据机器狗的状态（静止或移动）设置目标角度。
    参数:
    bac_statue_p_or_pr (int): 机器狗的状态标志，0表示静止，1表示移动。

    返回值:
    无
    """
    global Sta_Pitch, Sta_Roll
    import dog_data
    import dog_cfg
    get_imu_value()  # 获取IMU传感器的数据

    # 根据滤波后的俯仰角数据和陀螺仪数据调整Sta_Pitch
    Sta_Pitch = Sta_Pitch - (filter_data_p - p_origin) * kp_sta - gyro_x_fitted * kd_sta
    # 限制Sta_Pitch在允许的最大角度范围内
    if Sta_Pitch >= dog_cfg.pit_max_ang:
        Sta_Pitch = dog_cfg.pit_max_ang
    elif Sta_Pitch <= -dog_cfg.pit_max_ang:
        Sta_Pitch = -dog_cfg.pit_max_ang

    # 根据滤波后的横滚角数据调整Sta_Roll
    Sta_Roll = Sta_Roll - (filter_data_r - r_origin) * kp_sta
    # 限制Sta_Roll在允许的最大角度范围内
    if Sta_Roll >= dog_cfg.rol_max_ang:
        Sta_Roll = dog_cfg.rol_max_ang
    elif Sta_Roll <= -dog_cfg.rol_max_ang:
        Sta_Roll = -dog_cfg.rol_max_ang

    # 如果陀螺仪校准完成
    if gyro_cal_sta == 1:
        # 根据机器狗的状态设置目标角度
        if bac_statue_p_or_pr == 0:  # 机器狗静止
            dog_cfg.PIT_goal = Sta_Pitch
            dog_cfg.ROL_goal = Sta_Roll
        else:  # 机器狗移动
            dog_cfg.PIT_goal = 0  # 清零，防止运动时过度调节
            dog_cfg.ROL_goal = 0
